#include <ntifs.h>
#include <ntstrsafe.h>
#include <windef.h>

#include "CRegistry.h"

EXTERN_C
NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT pDriverObject,
    _In_ PUNICODE_STRING pRegisterPath
);

UNICODE_STRING g_ustrRegsitryPath;


VOID 
TestRegistryKeyCreate(
    _In_ PUNICODE_STRING pRegistryPath
    )
{
    DbgBreakPoint();

    NTSTATUS status;
    WCHAR wszKeyPath[MAX_PATH];
    RtlStringCbPrintfW(wszKeyPath, sizeof(wszKeyPath), L"%s\\%s", pRegistryPath->Buffer, L"Test");
    KdPrint(("[Debug] wszKeyPath=%ws\r\n", wszKeyPath));
    
    wk::CRegistryObject reg;
    status = reg.CreateKey(wszKeyPath);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is fail, status=%#X, wszKeyPath=%ws\r\n", status, wszKeyPath));
        return;
    }
    else
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is success, hKey=%#X, Path=%wZ\r\n", reg.GetKeyHandle(), reg.GetRegistryPath()));
    }
}

VOID 
TestRegistryKeyDelete(
    _In_ PUNICODE_STRING pRegistryPath
    )
{
    DbgBreakPoint();

    NTSTATUS status;
    WCHAR wszKeyPath[MAX_PATH];
    RtlStringCbPrintfW(wszKeyPath, sizeof(wszKeyPath), L"%s\\%s", pRegistryPath->Buffer, L"Test");
    KdPrint(("[Debug] wszKeyPath=%ws\r\n", wszKeyPath));

    wk::CRegistryObject reg;
    status = reg.OpenKey(wszKeyPath);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] CRegistryObject.OpenKey() is fail, status=%#X, wszKeyPath=%ws\r\n", status, wszKeyPath));
        return;
    }
    else
    {
        KdPrint(("[Debug] CRegistryObject.OpenKey() is success, hKey=%#X, Path=%wZ\r\n", reg.GetKeyHandle(), reg.GetRegistryPath()));
    }
    status = reg.DeleteKey();
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] CRegistryObject.DeleteKey() is fail, status=%#X, Registry=%wZ\r\n", status, reg.GetRegistryPath()));
        return;
    }
}

VOID 
TestRegistrySubKeyCreate(
    _In_ PUNICODE_STRING pRegistryPath
    )
{
    NTSTATUS status;
    WCHAR wszKeyPath[MAX_PATH];
    RtlStringCbPrintfW(wszKeyPath, sizeof(wszKeyPath), L"%s\\%s", pRegistryPath->Buffer, L"Test");
    KdPrint(("[Debug] wszKeyPath=%ws\r\n", wszKeyPath));

    wk::CRegistryObject reg;
    status = reg.CreateKey(wszKeyPath);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is fail, status=%#X, wszKeyPath=%ws\r\n", status, wszKeyPath));
        return;
    }
    else
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is success, hKey=%#X, Path=%wZ\r\n", reg.GetKeyHandle(), reg.GetRegistryPath()));
    }

    DbgBreakPoint();
    wk::CRegistryObject L1, L2;
    status = reg.CreateSubKey(L"1", &L1);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] Test Layer Add SubLayer is fail\r\n"));
    }
    else
    {
        KdPrint(("[Debug] Test Layer Add SubLayer is success\r\n"));
    }


    status = L1.CreateSubKey(L"A", &L2);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] 1 Layer Add SubLayer is fail\r\n"));
    }
    else
    {
        KdPrint(("[Debug] 1 Layer Add SubLayer is Success\r\n"));
    }
}

VOID
TestRegistrySubKeyDelete(
    _In_ PUNICODE_STRING pRegistryPath
    )
{
    NTSTATUS status;
    WCHAR wszKeyPath[MAX_PATH];
    RtlStringCbPrintfW(wszKeyPath, sizeof(wszKeyPath), L"%s\\%s", pRegistryPath->Buffer, L"Test");
    KdPrint(("[Debug] wszKeyPath=%ws\r\n", wszKeyPath));

    wk::CRegistryObject reg;
    status = reg.CreateKey(wszKeyPath);
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is fail, status=%#X, wszKeyPath=%ws\r\n", status, wszKeyPath));
        return;
    }
    else
    {
        KdPrint(("[Debug] CRegistryObject.CreateKey() is success, hKey=%#X, Path=%wZ\r\n", reg.GetKeyHandle(), reg.GetRegistryPath()));
    }

    DbgBreakPoint();
    wk::CRegistryObject subLayer;
    reg.OpenSubKey(L"1", &subLayer);

    status = subLayer.DeleteSubKey(L"A");
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] Delelte A Layer is fail, status=%#X\r\n", status));
    }
    else
    {
        KdPrint(("[Debug] Delelte A Layer is Success\r\n"));
    }

    status = reg.DeleteSubKey(L"1");
    if (!NT_SUCCESS(status))
    {
        KdPrint(("[Debug] Delelte 1 Layer is fail, status=%#X\r\n", status));
    }
    else
    {
        KdPrint(("[Debug] Delelte 1 Layer is Success\r\n"));
    }
}



VOID
DriverUnload(
    _In_ struct _DRIVER_OBJECT* DriverObject
)
{
    KdPrint(("[Debug] DriverUnload()\r\n"));

    UNREFERENCED_PARAMETER(DriverObject);

    //TestRegistryKeyDelete(&g_ustrRegsitryPath);
    TestRegistrySubKeyDelete(&g_ustrRegsitryPath);
    RtlFreeUnicodeString(&g_ustrRegsitryPath);
}

NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT pDriverObject,
    _In_ PUNICODE_STRING pRegistryPath
)
{
    KdPrint(("[Debug] DriverEntry()\r\n"));

    UNREFERENCED_PARAMETER(pDriverObject);

    pDriverObject->DriverUnload = DriverUnload;

    //TestRegistryKeyCreate(pRegistryPath);
    TestRegistrySubKeyCreate(pRegistryPath);
    
    RtlCreateUnicodeString(&g_ustrRegsitryPath, pRegistryPath->Buffer);


    return STATUS_SUCCESS;
}

